from .particle import Particle
from lib.utils.get_config import get_cfg
from lib.trainer import train
import os.path as osp
import pandas as pd
from tqdm.notebook import tqdm

cfg = get_cfg(osp.join(osp.abspath("./"), 'configs/config.yaml'))


class Swarm:
    def __init__(self, informants_number, particle_number):
        self.best_loss = 1000
        self.best_accuracy = 0
        self.informants_number = informants_number
        self.swarm = []
        self.particle_number = particle_number
        self.average_accuracy = 0
        self.average_loss = 0
        self.best_lr = 0
        self.best_hidden_num = 0
        self.best_hidden_dim = 0

        for _ in range(self.particle_number):
            new_particle = Particle()
            self.swarm.append(new_particle)

        for j in range(len(self.swarm)):
            self.swarm[j].setInformants(self.swarm, self.informants_number, j)

    def Optimize(self):
        # Run Optimization
        for p in range(self.particle_number):
            # Find best informants
            self.swarm[p].set_informant_best()
            # Update velocities
            self.swarm[p].update_velocity()
            # Apply velocities the position
            self.swarm[p].change_wb()

    def swarm_iteration(self, dataset, task):
        self.best_test_accuracy = 0
        self.average_test_accuracy = 0
        self.average_test_loss = 0
        for p in range(int(self.particle_number)):
            # informants_best[0] = hidden_num
            # informants_best[1] = hidden_dim
            # informants_best[2] = lr

            model, validation_loss, validation_accuracy, test_accuracy, test_loss = train(dataset, task, int(round(self.swarm[p].hidden_num)),
                                                                                          int(round(self.swarm[p].hidden_dim)), self.swarm[p].lr)

            if test_accuracy > self.best_test_accuracy:
                self.best_test_loss = test_loss
                self.best_test_accuracy = test_accuracy
                self.best_lr = self.swarm[p].lr
                self.best_hidden_num = int(round(self.swarm[p].hidden_num))
                self.best_hidden_dim = int(round(self.swarm[p].hidden_dim))

            self.average_test_accuracy += test_accuracy
            self.average_test_loss += test_loss

            self.swarm[p].check_error(validation_loss)

        self.average_test_accuracy = self.average_test_accuracy / \
            int(self.particle_number)
        self.average_test_loss = self.average_test_loss / \
            int(self.particle_number)

        return self.best_test_loss, self.best_test_accuracy, self.average_test_loss, self.average_test_accuracy, self.best_hidden_num, self.best_hidden_dim, self.best_lr

    def get_best(self):
        # Return the best error rate of the entire swarn
        swarm_best_err = -1
        for p in self.swarm:
            if p.best_err < swarm_best_err or swarm_best_err == -1:
                swarm_best_err = p.best_err

        return swarm_best_err

    def run(self, dataset, task):
        res_df = pd.DataFrame(data={
            'epoch':  [],
            'loss': [],
            'accuracy': [],
            'average_loss': [],
            'average_accuracy': [],
            'best_particle_hidden_num': [],
            'best_particle_hidden_dim': [],
            'best_particle_lr': []
        })
        # create a swarm with NUM_PARTICLES particles with NUM_INFORMANTS informants
        for epoch in tqdm(range(cfg['EPOCHS'])):
            self.swarm_iteration(dataset, task)
            self.Optimize()
            if epoch % cfg['PRINT_FREQ'] == 0:
                print(
                    f"Epoch = {epoch+1}/{cfg['EPOCHS']}\tLoss: {self.best_test_loss:.3f}\tAccuracy: {self.best_test_accuracy:.3f}")
            res_df = res_df.append({
                'epoch': epoch,
                'loss': self.best_test_loss,
                'accuracy': self.best_test_accuracy,
                'average_loss': self.average_test_loss,
                'average_accuracy': self.average_test_accuracy,
                'best_particle_hidden_num': self.best_hidden_num,
                'best_particle_hidden_dim': self.best_hidden_dim,
                'best_particle_lr': self.best_lr,
                'best_err': self.get_best()
            }, ignore_index=True)
        return res_df
